import torch
from time import time


@torch.no_grad()
def test(epoch_, loader, model, loss_func, device=torch.device("cpu")) -> tuple:
    run_loss = 0.0
    total = 0
    correct = 0
    model.eval()
    start_time = time()
    for X, Y in loader:
        X, Y = X.to(device), Y.to(device)

        pred_y = model(X)

        loss = loss_func(pred_y, Y)

        run_loss += loss.item()

        total += Y.size(0)
        pred_y = torch.max(pred_y, 1)[1]
        correct += (pred_y == Y).sum().item()
    end_time = time() - start_time
    print(f'Train {epoch_}: total = {total}, loss = {run_loss / total}, acc = {correct / total}, using time = {end_time}s')
    return run_loss / total, correct / total
